#include "Integrator.h"
#include "FastOps.h"
#include "FileIO.h"
#include "Vector.h"
#include "SystemData.h"
#include "MultibodySystem.h"
#include "State.h"
#include "Reporter.h"

#include <math.h>
#include <cstdlib>
#include <assert.h>
#include <fstream>
#include <iomanip>

#define max(x,y) ( (x) < (y) ? (y) : (x) )
#define min(x,y) ( (x) < (y) ? (x) : (y) )

#ifndef RSIM_SMALL
        #define RSIM_SMALL 1.0E-10
#endif

using namespace RSIM;
using namespace std;

///////////////////////////////////////////////////////	

Integrator::Integrator():	hmax_(1.0),hmin_(1.0E-15),RepObj_(NULL),numY_(0),
                                AbsTol_(1.0E-7),RelTol_(1.0E-7),Threshold_(1.0E-3),
                                MaxNumCuts_(20),MaxScaleFactor_(5),MinScaleFactor_(0.1),
                                Refine_(4){}

///////////////////////////////////////////////////////	

void Integrator::adjustStepSize(const double& error, 
                                const int& ErrorOrder,
                                const double& h,
                                bool& noFailInThisStep,
                                double& h_new){

        // See \cite{hairer1993} page 168, eqns 4.11-4.13
        const double fac = 0.8;
        
        if(error==0.0){
                // If the error is 0, scale h by the greatest value. 
                h_new = MaxScaleFactor_*h;
                return;
        }		
        
        if(noFailInThisStep){
                if(error<RelTol_){
                        // h can at most be scaled by MaxScaleFactor_
                        h_new = h*min(fac/pow((error/RelTol_),1.0/double(ErrorOrder)),MaxScaleFactor_);
                        return;
                }
        }
        
        if(error>=RelTol_){
                if(noFailInThisStep){
                        // first cut in time-step is conservative
                        noFailInThisStep = false;
                        h_new = h*max(MinScaleFactor_, fac*pow(RelTol_/error, 1.0/ErrorOrder));
                }
                else{
                        // Second cut onwards are all by a factor of 0.5
                        h_new = h*0.5;
                }
                if(h_new < hmin_){
                        cout<<"Time step required to maintain tolerances is less than the specified minimum. "
                        <<"Change using setHMin()\n";
                        exit(0);
                }
                return;
        }
}

///////////////////////////////////////////////////////

void Integrator::stepTo(void(*calcYDot)(const double& ,	// Time (t)
                                        const Vector&, 	// Vector Y(t)
                                        Vector&  	// Vector YDot(t)
                                        ),
                        const double& t_initial,
                        const double& t_final,
                        const Vector& Y0,
                        const int& writeStateVectorAtEachTimeStep,
                        const char FileName[],
                        const double& deltaT
                        ){
        
        ////////////////////////////////////////////////////////
        ////////////////////////////////////////////////////////
        // Initialization of the Integrator
        assert(t_initial<t_final);
        assert(deltaT > 0 );
        this->numY_ = Y0.getLength();
        
        // set the size of YDot at t
        Vector Y0Dot(numY_);
        
        // set the size of YDot at t
        Vector Y1Dot(numY_);
        
        // calculate YDot at time t_initial
        // Useful for integrators which belong to the class First Same As Last (FSAL)
        calcYDot(t_initial, Y0, Y0Dot);

        // if not implemented in a derieved class, this does nothing
        this->initIntegrator(Y0);
        
        // stores the intermediate times at which interpolated values are to be computed
        double *t_interp=NULL; 
        
        // stores the interpolated values
        Vector y_interp; 
        
        // Output written in this file 
        fstream OpFile;
        
        if(writeStateVectorAtEachTimeStep){		
                // Open file in which state vector will be written        
                OpFile.open(FileName,ios::app | ios::out );
                
                if(Refine_>1){
                        y_interp.resize(numY_);
                        t_interp = new double[Refine_-1];
                        for(int i=0,j=1; i<Refine_-1; t_interp[i++] = double(j++)/double(Refine_));
                }
        }
                
        // this value is set by the derived class.
        const int error_order = this->getErrorOrder();
        
        // h is an esitmate of the first time step
        double h = estimateInitialTimeStep(t_initial, t_final, error_order, Y0, Y0Dot);
        
        double h_new;
        
        double current_time = t_initial;
        
        // y0 stores the state vector before taking the step
        Vector y0 = Y0;
        
        // y1 stores state vector after taking the step
        Vector y1(numY_);
        
        // local error and tolerance in each time step 
        double step_error;
        
        // Initialization compplete here
        ////////////////////////////////////////////////////////
        ////////////////////////////////////////////////////////
        
        // step_count: number of times a time-step can be cut down.
        // numSteps: counts the total number of time steps taken.
        int step_count, numSteps=0;
        while(1){
                bool noFailInThisStep = true;
                for(step_count=0;step_count<MaxNumCuts_;++step_count){

                        // Take a tentative step. calculate the step_error
                        this->attemptSingleStep(calcYDot, current_time, y0,Y0Dot,  h, y1,Y1Dot, step_error);

                        // adjust step size based on the step error and tolerance
                        this->adjustStepSize(step_error, error_order, h, noFailInThisStep, h_new);

                        // accept step if step_error is lower than tolerance
                        if(step_error<=RelTol_){
                                break;
                        }
                        else{
                                h = h_new;
                        }
                }

                // if the step size is cut more than MaxNumCuts, exit the code.
                if(step_count>=MaxNumCuts_){
                        cout<<"\nIntegrator::stepTo(), Integration failed to converge.\n"<<
                        "The lowest time step used was "<<h<<endl; 
                        exit(0);
                }		                               
                ++numSteps;
                if(writeStateVectorAtEachTimeStep){
                        FileIO::write_append_SingleArrayToFile(current_time, y0.Ptr(), numY_,FileName);
                        // interpolate if required to compute intermediate values
                        for(int i=0;i<Refine_-1;++i){
                                interpolate(current_time,y0,Y0Dot,Y1Dot, h,t_interp[i],y_interp);				
                                FileIO::write_append_SingleArrayToFile(OpFile, current_time+t_interp[i]*h, y_interp.Ptr(), numY_, FileName);
                        }
                }
                                
                // increment current time 
                current_time += h;
                
                // break if integration is complete
                if(current_time==t_final){
                        if(writeStateVectorAtEachTimeStep)
                                FileIO::write_append_SingleArrayToFile(OpFile,current_time, y1.Ptr(), numY_,FileName);
                        break;
                }						
                
                // hit the endpoint exaclty and don't allow h to grow beyond used specified value
                h = min(min(h_new,this->hmax_),t_final-current_time);
                
                y0 = y1;
                Y0Dot = Y1Dot;
        } // while(1) // this is the main loop
        
        // t_interp is not needed anymore.
        delete [] t_interp;
        OpFile.close();
}

///////////////////////////////////////////////////////	

void Integrator::stepTo(const double& t_initial,
                        const double& t_final,
                        MultibodySystem& MBS, 
                        State& state,
                        const int& writeStateVectorAtEachTimeStep,
                        const char FileName[],
                        const double& deltaT){
        
        ////////////////////////////////////////////////////////
        ////////////////////////////////////////////////////////
        // Initialization of the Integrator
        assert(t_initial<t_final);
        assert(deltaT > 0 );
        
        state.setTime(t_initial);
        
        // We are integrating both the Q's and U's of the system.
        // Thus size of the state vector is y0 = [q,u]'
        const int& nq = state.getNQ();
        const int& nu = state.getNU();
        
        this->numY_ = nq + nu;
        
        Vector Y0(numY_);
        
        state.getY(Y0.wPtr());
        
        // set the size of YDot at t
        Vector Y0Dot(numY_);
        
        // set the size of YDot at t
        Vector Y1Dot(numY_);
        
        // calculate YDot at time t_initial
        // Useful for integrators which belong to the class First Same As Last (FSAL)
        MBS.calcSystemAcceleration(state);	
        state.getYDot(Y0Dot.wPtr());	
        
        // if not implemented in a derieved class, this does nothing
        this->initIntegrator(Y0);
        
        // stores the intermediate times at which interpolated values are to be computed
        double *t_interp=NULL; 
        
        // stores the interpolated values
        Vector y_interp; 
        
        // Output written in this file 
        fstream OpFile;
        
        if(writeStateVectorAtEachTimeStep){		
                // Open file in which state vector will be written
                OpFile.open(FileName,ios::app | ios::out );
                
                if(Refine_>1){
                        y_interp.resize(numY_);
                        t_interp = new double[Refine_-1];
                        for(int i=0,j=1; i<Refine_-1; t_interp[i++] = double(j++)/double(Refine_));
                }
        }
        
        // this value is set by the derived class.
        const int error_order = this->getErrorOrder();
                
        // h is an esitmate of the first time step
        double h = estimateInitialTimeStep(t_initial, t_final, error_order, Y0, Y0Dot);
        
        double h_new;
        
        double current_time = t_initial;		
        
        // y1 stores state vector after taking the step
        Vector Y1(numY_);
        
        // local error and tolerance in each time step 
        double step_error;
        
        // Initialization compplete here
        ////////////////////////////////////////////////////////
        ////////////////////////////////////////////////////////
        
        // step_count: number of times a time-step can be cut down.
        // numSteps: counts the total number of time steps taken.
        int step_count, numSteps=0;

        while(1){
                bool noFailInThisStep = true;
                for(step_count=0;step_count<MaxNumCuts_;++step_count){

                        // Take a tentative step. calculate the step_error
                        this->attemptSingleStep(current_time, Y0,Y0Dot, h,MBS, state, Y1,Y1Dot, step_error);

                        // adjust step size based on the step error and tolerance
                        this->adjustStepSize(step_error, error_order, h, noFailInThisStep, h_new);

                        // accept step if step_error is lower than tolerance
                        if(step_error<=RelTol_){
                                break;
                        }
                        else{
                                h = h_new;
                        }
                }

                // if the step size is cut more than MaxNumCuts, exit the code.
                if(step_count>=MaxNumCuts_){
                        cout<<"\nIntegrator::stepTo(), Integration failed to converge.\n"<<
                        "The lowest time step used was "<<h<<endl; 
                        exit(0);
                }		
                ++numSteps;
                if(writeStateVectorAtEachTimeStep){
                        FileIO::write_append_SingleArrayToFile(current_time, Y0.Ptr(), numY_,FileName);
                        
                        // interpolate if required to compute intermediate values
                        for(int i=0;i<Refine_-1;++i){
                                interpolate(current_time,Y0,Y0Dot,Y1Dot, h,t_interp[i],y_interp);				
                                FileIO::write_append_SingleArrayToFile(OpFile, current_time+t_interp[i]*h, y_interp.Ptr(), numY_, FileName);
                        }
                }
                                
                // increment current time 
                current_time += h;
                
                // break if integration is complete
                if(current_time==t_final){
                        if(writeStateVectorAtEachTimeStep)
                                FileIO::write_append_SingleArrayToFile(OpFile,current_time, Y1.Ptr(), numY_,FileName);
                        break;
                }						

                // hit the endpoint exaclty and don't allow h to grow beyond used specified value
                h = min(min(h_new,this->hmax_),t_final-current_time);

                Y0 = Y1;
                Y0Dot = Y1Dot;
        } // while(1) // this is the main loop
        
        // t_interp is not needed anymore.
        delete [] t_interp;
        OpFile.close();

}

////////////////////////////////////////////////////////

void Integrator::extractYFromState(const SystemData& Data,
                                   const Vector_<int>& NY,
                                   Vector& Y)const{
        int tmp_var = 0;
        double *ywptr = Y.wPtr();
        for(int i=0;i<Data.system_.getLength();++i){
                Data.state_[i]->getY(ywptr+tmp_var);
                tmp_var += NY[i];
        }
}

////////////////////////////////////////////////////////

void Integrator::extractYDotFromState(  const SystemData& Data,
                                        const Vector_<int>& NY,
                                        Vector& YDot)const{
        int tmp_var = 0;
        double *ydot = YDot.wPtr();
        for(int i=0;i<Data.system_.getLength();++i){
                Data.state_[i]->getYDot(ydot+tmp_var);
                tmp_var += NY[i];
        }
}

////////////////////////////////////////////////////////                        

void Integrator::calcSystemAcceleration(SystemData& Data){
        for(int i=0;i<Data.system_.getLength();++i){
                State& state = *Data.state_[i];
                MultibodySystem& MBS = *Data.system_[i];
                MBS.calcSystemAcceleration(state);
        }        
}
                        
////////////////////////////////////////////////////////                        

void Integrator::stepTo(const double& t_final,
                        SystemData& Data){                                
        double report_interval = 0.0;        
        assert(Data.Time()<t_final);
        
        ////////////////////////////////////////////////////////
        ////////////////////////////////////////////////////////
        // Initialization of the Integrator starts here.        
        
        // Total generalized coordinates and speeds active in the all the systems
        numY_ = Data.getNY();
        
        Vector Y0(numY_);
        
        // stores the interpolated values
        Vector y_interp;
        
        if(RepObj_)
                y_interp.resize(numY_);
        
        // Get the initial state of the system
        Data.extractYFromState(Y0);        
                                
        // set the size of YDot at t
        Vector Y0Dot(numY_);
        
        // set the size of YDot at t
        Vector Y1Dot(numY_);
        
        // calculate YDot at time t_initial
        // Useful for integrators which belong to the class First Same As Last (FSAL)
        Data.calcSystemAcceleration();

        Data.extractYDotFromState(Y0Dot);

        if(RepObj_){
                report_interval = RepObj_->getReportInterval();
                
                assert(report_interval>0);
                assert(report_interval < (t_final-Data.Time()));
                
                // Report is called by default at time t_initial and t_final
                // this is the call for t_initial. It is assumed that the state
                // provided is consistent for t_initial
                RepObj_->report();              
        }    
        
        // if not implemented in a derieved class, this does nothing
        initIntegrator(Y0);                               
        
        // this value is set by the derived class.
        const int error_order = this->getErrorOrder();
        
        // h is an esitmate of the first time step. Very conservative.
        double h = estimateInitialTimeStep(Data.Time(), t_final, error_order, Y0, Y0Dot);
        
        double h_new;
        
        double current_time = Data.Time();             
        
        // Calculate time for next report
        double ReportNextAtTime = current_time+report_interval;
        
        // y1 stores state vector after taking the step
        Vector Y1(numY_);
        
        // local error and tolerance in each time step 
        double step_error;
        
        // Initialization compplete here
        ////////////////////////////////////////////////////////
        ////////////////////////////////////////////////////////
        
        // step_count: number of times a time-step can be cut down.
        // numSteps: counts the total number of time steps taken.
        int step_count, numSteps=0;                     
        
        while(1){
                bool noFailInThisStep = true;
                for(step_count=0;step_count<MaxNumCuts_;++step_count){
                        
                        // Take a tentative step. calculate the step_error
                        this->attemptSingleStep(current_time, 
                                                Y0,
                                                Y0Dot, 
                                                h,
                                                Data, 
                                                Y1,
                                                Y1Dot, 
                                                step_error);
                        
                        // adjust step size based on the step error and tolerance
                        this->adjustStepSize(step_error, 
                                             error_order, 
                                             h, 
                                             noFailInThisStep, 
                                             h_new);
                        
                        // accept step if step_error is lower than tolerance
                        if(step_error<=RelTol_){
                                break;
                        }
                        else{
                                h = h_new;
                        }
                } 
                
                // if the step size is cut more than MaxNumCuts, exit the code.
                if(step_count>=MaxNumCuts_){
                        cout<<"\nIntegrator::stepTo(), Integration failed to converge.\n"<<
                        "The lowest time step used was "<<h<<endl; 
                        exit(0);
                }
                ++numSteps;
                
                if(RepObj_ ){
                        for(;ReportNextAtTime<=(current_time+h+RSIM_SMALL);){
                                
                                // No interpolation required if we are too close to the current_time_step
                                // or current_time + h. So we directly report and proceed.
                                if(fabs(ReportNextAtTime-current_time)<RSIM_SMALL){
                                        Data.setY(Y0);
                                        Data.Time()=current_time;
                                        RepObj_->report();
                                        ReportNextAtTime += report_interval;
                                        continue;
                                }
                                
                                if(fabs(ReportNextAtTime-current_time-h)<RSIM_SMALL){
                                        Data.setY(Y1);
                                        Data.Time()=current_time+h;
                                        RepObj_->report();
                                        ReportNextAtTime += report_interval;
                                        continue;
                                }
                                
                                // interpolate to compute intermediate values
                                interpolate(current_time, 
                                            Y0, 
                                            Y0Dot, 
                                            Y1Dot, 
                                            h, 
                                            (ReportNextAtTime-current_time)/h, 
                                            y_interp);
                                
                                // set the intermediate state vector
                                Data.setY(y_interp);
                                Data.Time() = ReportNextAtTime;
                                
                                // Report
                                RepObj_->report();
                                
                                ReportNextAtTime += report_interval;
                        }
                } // if(RepObj_)
                
                // increment current time 
                current_time += h;
                
                // break if integration is complete
                if(current_time==t_final){
                        break;
                }
                
                // hit the endpoint exaclty and don't allow h to grow beyond user specified value
                h = min(min(h_new,this->hmax_),t_final-current_time);
                
                Y0 = Y1;
                Y0Dot = Y1Dot;
        } // while(1) // this is the main loop          
}
                        
///////////////////////////////////////////////////////

void Integrator::stepTo(const double& t_final,
                        MultibodySystem& MBS, 
                        State& state){

        ////////////////////////////////////////////////////////
        ////////////////////////////////////////////////////////

        double report_interval = 0.0;
                
        // Initialization of the Integrator
        assert(state.getTime()<t_final);        
                        
        // We are integrating both the Q's and U's of the system.
        // Thus size of the state vector is y0 = [q,u]'
        const int& nq = state.getNQ();
        const int& nu = state.getNU();
        
        this->numY_ = nq + nu;
        
        Vector Y0(numY_);
        
        // stores the interpolated values
        Vector y_interp;
        
        if(RepObj_)
                y_interp.resize(numY_);
        
        state.getY(Y0.wPtr());
        
        // set the size of YDot at t
        Vector Y0Dot(numY_);
        
        // set the size of YDot at t
        Vector Y1Dot(numY_);
        
        // calculate YDot at time t_initial
        // Useful for integrators which belong to the class First Same As Last (FSAL)
        MBS.calcSystemAcceleration(state);	
        state.getYDot(Y0Dot.wPtr());
        
        if(RepObj_){
                report_interval = RepObj_->getReportInterval();
                
                assert(report_interval>0);
                assert(report_interval < (t_final-state.getTime()));
                
                // Report is called by default at time t_initial and t_final
                // this is the call for t_initial. It is assumed that the state
                // provided is consistent for t_initial
                RepObj_->report();              
        }               
        
        // if not implemented in a derieved class, this does nothing
        this->initIntegrator(Y0);				
                
        // this value is set by the derived class.
        const int error_order = this->getErrorOrder();
                
        // h is an esitmate of the first time step
        double h = estimateInitialTimeStep(state.getTime(), t_final, error_order, Y0, Y0Dot);
        
        double h_new;
        
        double current_time = state.getTime();
        
        // Calculate time for next report
        double ReportNextAtTime = current_time+report_interval;
        
        // y1 stores state vector after taking the step
        Vector Y1(numY_);
        
        // local error and tolerance in each time step 
        double step_error;
        
        // Initialization compplete here
        ////////////////////////////////////////////////////////
        ////////////////////////////////////////////////////////
        
        // step_count: number of times a time-step can be cut down.
        // numSteps: counts the total number of time steps taken.
        int step_count, numSteps=0;			
        
        while(1){
                bool noFailInThisStep = true;
                for(step_count=0;step_count<MaxNumCuts_;++step_count){

                        // Take a tentative step. calculate the step_error
                        this->attemptSingleStep(current_time, Y0,Y0Dot, h,MBS, state, Y1,Y1Dot, step_error);

                        // adjust step size based on the step error and tolerance
                        this->adjustStepSize(step_error, error_order, h, noFailInThisStep, h_new);

                        // accept step if step_error is lower than tolerance
                        if(step_error<=RelTol_){
                                break;
                        }
                        else{
                                h = h_new;
                        }
                } 

                // if the step size is cut more than MaxNumCuts, exit the code.
                if(step_count>=MaxNumCuts_){
                        cout<<"\nIntegrator::stepTo(), Integration failed to converge.\n"<<
                        "The lowest time step used was "<<h<<endl; 
                        exit(0);
                }
                ++numSteps;

                if(RepObj_ ){
                        for(;ReportNextAtTime<=(current_time+h+RSIM_SMALL);){
                                
                                // No interpolation required if we are too close to the current_time_step
                                // or current_time + h. So we directly report and proceed.
                                if(fabs(ReportNextAtTime-current_time)<RSIM_SMALL){
                                        state.setY(Y0.Ptr());
                                        state.setTime(current_time);
                                        RepObj_->report();
                                        ReportNextAtTime += report_interval;
                                        continue;
                                }
                                
                                if(fabs(ReportNextAtTime-current_time-h)<RSIM_SMALL){
                                        state.setY(Y1.Ptr());
                                        state.setTime(current_time+h);
                                        RepObj_->report();
                                        ReportNextAtTime += report_interval;
                                        continue;
                                }

                                // interpolate to compute intermediate values
                                interpolate(current_time, Y0, Y0Dot, Y1Dot, h, (ReportNextAtTime-current_time)/h, y_interp);
                                
                                // set the intermediate state vector
                                state.setY(y_interp.Ptr());
                                state.setTime(ReportNextAtTime);

                                // Report
                                RepObj_->report();

                                ReportNextAtTime += report_interval;
                        }
                }

                // increment current time 
                current_time += h;

                // break if integration is complete
                if(current_time==t_final){
                        break;
                }

                // hit the endpoint exaclty and don't allow h to grow beyond user specified value
                h = min(min(h_new,this->hmax_),t_final-current_time);

                Y0 = Y1;
                Y0Dot = Y1Dot;
        } // while(1) // this is the main loop		
}

///////////////////////////////////////////////////////

void Integrator::estimateLocalStepError(const Vector& Y0, 
                                        const Vector& Y1, 
                                        const Vector& Yerr, 
                                        const double& current_time_step, 
                                        double& error){
        const double *y0 = Y0.Ptr();
        const double *y1 = Y1.Ptr();
        const double *yerr = Yerr.Ptr();
        error = 0.0;
        
        for(int i=0;i<numY_;++i){
                error = max(error, fabs(yerr[i])/max(max(fabs(y0[i]),fabs(y1[i])),Threshold_));
        }
        error *= current_time_step;
}

///////////////////////////////////////////////////////

double Integrator::estimateInitialTimeStep(const double& t_initial, 
                                        const double& t_final, 
                                        const int& error_order,
                                        const Vector& Y, 
                                        const Vector& YDot){
        const double *y = Y.Ptr();
        const double *ydot = YDot.Ptr();	
        double h_tmp = 0.0;
        double h;
        h = min(t_final-t_initial, 2.0);
        for(int i=0;i<numY_;++i){
                h_tmp = max(h_tmp, fabs(ydot[i])/max(fabs(y[i]),Threshold_));
        }

        h_tmp /= 0.8*pow(RelTol_,1.0/double(error_order));		
        if(h*h_tmp>1){
                h = 1/h_tmp;
                if(h<1E-15){
                        cout<<"Initial time-step too low\n";
                        exit(0);
                }
        }
        return h;
}

///////////////////////////////////////////////////////	

Integrator::~Integrator(){
        delete RepObj_;
}

///////////////////////////////////////////////////////	

// @book{hairer1993,
//   title={Solving ordinary differential equations: Nonstiff problems},
//   author={Hairer, E. and N{\o}rsett, S.P. and Wanner, G.},
//   year={1993},
//   publisher={Springer}
// }

// @book{ascher1998,
//   title={{Computer methods for ordinary differential equations and differential-algebraic equations}},
//   author={Ascher, U.M. and Petzold, L.R.},
//   year={1998},
//   publisher={Society for Industrial Mathematics}
// }

///////////////////////////////////////////////////////	